import timm
import numpy as np
import torch
import random


def load_checkpoint(path):
    new_state_dict = {}
    state_dict = torch.load(path)
    for k, v in state_dict.items():
        new_state_dict[k.replace('module.', '')] = v
    return new_state_dict


def get_model(model_name, pretrained=False):
    model = timm.create_model(model_name, pretrained)
    return model


def random_seed(seed_value, use_cuda):
    np.random.seed(seed_value) # cpu vars
    torch.manual_seed(seed_value) # cpu  vars
    random.seed(seed_value) # Python
    if use_cuda:
        torch.cuda.manual_seed(seed_value)
        torch.cuda.manual_seed_all(seed_value) # gpu vars
        torch.backends.cudnn.deterministic = True  #needed
        torch.backends.cudnn.benchmark = True


def tuple_tensor_to_float(tu):
    return [x.float() for x in tu]
